package io.gitee.luoyuan800.tool.llm.agent;

import io.gitee.luoyuan800.tool.llm.bean.completion.ChatCompletion;
import io.gitee.luoyuan800.tool.llm.bean.config.AiAgentConfig;
import io.gitee.luoyuan800.tool.llm.bean.request.BaiduAiChatRequest;
import io.gitee.luoyuan800.tool.llm.bean.request.ChatRequest;
import io.gitee.luoyuan800.tool.llm.web.RestfulClient;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.util.Map;
import java.util.stream.Collectors;

/**
 * 百度 AI 开发平台大模型API调用Agent的实现。
 * 可以调用星河大模型。
 * 特别针对开启search功能时候的返回消息处理。
 */
@Slf4j
public class BaiduAiStudioAgent extends CommonAgent {
    private static final String URL = "https://aistudio.baidu.com/llm/lmapi/v3/chat/completions";

    public BaiduAiStudioAgent(AiAgentConfig aiAgentConfig, RestfulClient restTemplate, String name) {
        super(aiAgentConfig, restTemplate, name);
        if (StringUtils.isBlank(aiAgentConfig.getUrl())) {
            aiAgentConfig.getModelConfig().setFullUrl(URL);
        }
    }

    @Override
    public String chatWithImage(String msg, String imageUrl, boolean enableThinking) {
        ChatCompletion<?> response = fullChat(msg, 1, imageUrl, enableThinking);
        if (response != null && response.getChoices() != null && !response.getChoices().isEmpty()) {
            String answer = buildMessageString(true, response.getChoices().get(0).getMessage());
            if (response.getSearchResults() != null && !response.getSearchResults().isEmpty()) {
                answer += "\n\n" + response.getSearchResults().stream()
                        .map(rs -> rs.getTitle() + ": " + rs.getUrl()).collect(Collectors.joining("\n"));
            }
            return answer;
        } else {
            return "";
        }
    }

    public String chat(String msg, boolean showThink) {
        ChatCompletion<?> response = fullChat(msg, 1, null, showThink);
        if (response != null && response.getChoices() != null && !response.getChoices().isEmpty()) {
            String answer = buildMessageString(showThink, response.getChoices().get(0).getMessage());
            if (response.getSearchResults() != null && !response.getSearchResults().isEmpty()) {
                answer += "\n\n" + response.getSearchResults().stream()
                        .map(rs -> rs.getTitle() + ": " + rs.getUrl()).collect(Collectors.joining("\n"));
            }
            return answer;
        } else {
            return "";
        }
    }

    @Override
    ChatRequest buildRequest(int resultCount) {
        BaiduAiChatRequest request = new BaiduAiChatRequest();
        setRequestParams(resultCount, request);
        if (getAiAgentConfig().isWebSearchCitation() || getAiAgentConfig().isWebSearchTrace()) {
            Map<String, Boolean> webSearch = Map.of(
                    "enable", true,
                    "enable_citation", getAiAgentConfig().isWebSearchCitation(),
                    "enable_trace", getAiAgentConfig().isWebSearchTrace());
            request.setWebSearch(webSearch);
        }
        return request;
    }

}
